"""
document_processor.py是EduRAG系统的核心模块之一，用于文档解析。
主要负责加载多种格式的文档（如.txt、.pdf等），并对其进行分层切分，生成父块和子块，为后续的向量存储和检索做好准备。

"""
# 导包
import os
import sys
from datetime import datetime

from langchain_community.document_loaders import TextLoader, UnstructuredMarkdownLoader
from langchain_text_splitters import MarkdownTextSplitter, CharacterTextSplitter

from base.logger import get_logger
from base.config import get_config
from rag_qa.edu_document_loaders import OCRPDFLoader, OCRIMGLoader, OCRDOCLoader, OCRPPTLoader

# 定义支持的文档格式（文件扩展名）
document_loads = {
    # 文本文件使用 TextLoader
    ".txt": TextLoader,
    # PDF 文件使用 OCRPDFLoader
    ".pdf": OCRPDFLoader,
    # Word 文件使用 OCRDOCLoader
    ".docx": OCRDOCLoader,
    # PPT 文件使用 OCRPPTLoader
    ".ppt": OCRPPTLoader,
    # PPTX 文件使用 OCRPPTLoader
    ".pptx": OCRPPTLoader,
    # JPG 文件使用 OCRIMGLoader
    ".jpg": OCRIMGLoader,
    # PNG 文件使用 OCRIMGLoader
    ".png": OCRIMGLoader,
    # Markdown 文件使用 UnstructuredMarkdownLoader
    ".md": UnstructuredMarkdownLoader
}

# 对base模块下的logger进行初始化和config的加载
logger = get_logger()
config = get_config()

"""
需求：从指定的文件夹加载所有支持的文档格式并添加元数据（如文件名、文件路径等）
思路：
    1：遍历指定文件夹下的所有文件
    2：根据文件扩展名判断文件类型
    3：根据文件类型选择对应的Loader进行加载（如果文件类型是.txt格式，需要指定编码方式utf-8）
    4：为加载的文档添加元数据（如学科名、文件路径、时间戳等）
"""
def load_documents_from_directory(folder_path: str) -> list:
    # 定义文档列表
    documents = []
    # 获取到目前支持的文档格式
    supported_extensions = document_loads.keys()

    # 获取学科名称
    subject_name = os.path.basename(folder_path).replace("_data", "")

    # 1：遍历指定文件夹下的所有文件
    for root, dirs, files in os.walk(folder_path):
        # 不考虑嵌套子文件夹的问题，只对根目录下的文件列表进行遍历
        for file in files:
            # 获取文件的绝对路径
            file_path = os.path.join(root, file)
            # 获取文件的扩展名
            extension_name = os.path.splitext(file_path)[1].lower()

            # 2：根据文件扩展名判断文件类型
            if extension_name in supported_extensions:
                try:
                    # 根据文件类型选择对应的Loader进行加载
                    loader_class = document_loads[extension_name]
                    # 3：根据文件类型选择对应的Loader进行加载（如果文件类型是.txt格式，需要指定编码方式utf - 8）
                    if extension_name == ".txt":
                        loader = loader_class(file_path, encoding="utf-8")
                    else:
                        loader = loader_class(file_path)

                    # 调用加载器的load方法加载文档
                    loaded_docs = loader.load()

                    # 4：为加载的文档添加元数据（如学科名、文件路径、时间戳等）
                    for doc in loaded_docs:
                        doc.metadata["source"] = subject_name
                        doc.metadata["file_path"] = file_path
                        doc.metadata["extension"] = extension_name
                        doc.metadata["timestamp"] = datetime.now().isoformat()

                    # 5：将加载的文档添加到文档列表中
                    documents.extend(loaded_docs)
                    logger.info(f"文档加载完成：{file_path}")
                except Exception as e:
                    logger.error(f"文档加载失败：{file_path}")
                    logger.error(e)
            else:
                logger.warning(f"不支持的文件格式：{file_path}")

    # 返回加载完成的文档列表
    return documents

"""
需求：处理文档并进行分层切分，返回切分后的文档块列表
思路：
1：获取所有加载完成的文档
2：初始化文档切分器（CharacterTextSplitter），设置参数（如分隔符、块大小、重叠大小等）
    2.1：对于markdown格式的文件，使用MarkdownTextSplitter
    2.2：对于其他格式的文件，使用CharacterTextSplitter
3：对每个文档进行切分，生成父块和子块
    3.1：将一个文档切分成多个块chunk（父块和子块）
    3.2：把每一个字块进一步切分成更小的子块（如句子级别sub_chunk）
    3.3：给每个子块添加元数据，包括父块的元数据和子块的索引（如chunk_index、sub_chunk_index等）
"""
def process_documents(directory_path, parent_chunk_size=config.PARENT_CHUNK_SIZE,
                      sub_chunk_size=config.CHILD_CHUNK_SIZE,
                      chunk_overlap=config.CHUNK_OVERLAP):
    # 1：获取所有加载完成的文档
    docs = load_documents_from_directory(directory_path)
    logger.info(f"加载完成{len(docs)}个文档")

    # 2：初始化文档切分器（CharacterTextSplitter），设置参数（如分隔符、块大小、重叠大小等）
    #   2.1：对于markdown格式的文件，使用MarkdownTextSplitter
    markdown_parent_splitter = MarkdownTextSplitter(chunk_size=parent_chunk_size, chunk_overlap=chunk_overlap)
    markdown_child_splitter = MarkdownTextSplitter(chunk_size=sub_chunk_size, chunk_overlap=chunk_overlap)

    #   2.2：对于其他格式的文件，使用CharacterTextSplitter
    other_parent_splitter = CharacterTextSplitter(chunk_size=parent_chunk_size, chunk_overlap=chunk_overlap)
    other_child_splitter = CharacterTextSplitter(chunk_size=sub_chunk_size, chunk_overlap=chunk_overlap)

    # 定义所有字块的列表
    child_chunks = []

    # 3：对每个文档进行切分，生成父块和子块
    for idx, doc in enumerate(docs):
        # 获取到文档的元数据的扩展名
        extension_name = doc.metadata["extension"]
        # 判断文档的格式是否是markdown格式
        is_markdown = (extension_name == ".md")

        # 如果is_markdown为True，则使用markdown_parent_splitter进行切分，否则使用other_parent_splitter进行切
        parent_splitter = markdown_parent_splitter if is_markdown else other_parent_splitter
        child_splitter = markdown_child_splitter if is_markdown else other_child_splitter

        logger.info(f"正在处理文档：{doc.metadata['file_path']}, 文档格式：{extension_name}，使用的切分器为：{parent_splitter.__class__.__name__}")

        # 3.1：将一个文档切分成多个块chunk（父块和子块）
        parent_docs = parent_splitter.split_documents([doc])
        logger.info(f"文档切分成{len(parent_docs)}个父块")

        for j, parent_doc in enumerate(parent_docs):
            # 3.2：把每一个字块进一步切分成更小的子块（如句子级别sub_chunk）
            # 生成一个父块唯一id，格式为：doc_idx_parent_j
            parent_id = f"doc_{idx}_parent_{j}"
            # 将父块id添加到父块的元数据中
            parent_doc.metadata["parent_id"] = parent_id
            # 将父块内容存储到元数据 中
            parent_doc.metadata["parent_content"] = parent_doc.page_content

            # 使用子块分词器将父块内容进行切分成多个子块
            child_docs = child_splitter.split_documents([parent_doc])
            logger.info(f"父块切分成{len(child_docs)}个子块")

            # 遍历每个子块，添加元数据
            for k, child_doc in enumerate(child_docs):
                # 3.3：给每个子块添加元数据，包括父块的元数据和子块的索引（如chunk_index、sub_chunk_index等）
                # 生成一个子块唯一id，格式为：parent_id_sub_k
                child_id = f"{parent_id}_child_{k}"
                # 将父块id添加到子块的元数据中
                child_doc.metadata["parent_id"] = parent_id
                # 将父块内容存储到子块元数据中
                child_doc.metadata["parent_content"] = parent_doc.page_content
                # 将子块id添加到子块的元数据中
                child_doc.metadata["id"] = child_id
                # 将子块添加到字块列表中
                child_chunks.append(child_doc)

    # 记录子块总数
    logger.info(f"子块总数：{len(child_chunks)}")

    # 返回所有字块列表
    return child_chunks

if __name__ == '__main__':
    print(process_documents("F://code//integrated_qa_system//rag_qa//data//ai_data"))
